from elasticsearch import Elasticsearch
from pymongo import MongoClient
from django.conf import settings
from .models import TPADocument


class DataSyncService:
    def __init__(self):
        self.mongo_client = MongoClient(settings.MONGODB_URI)
        self.db = self.mongo_client[settings.MONGODB_DB]
        self.collection = self.db[settings.MONGODB_COLLECTION]
        self.es = Elasticsearch([settings.ELASTICSEARCH_HOST],
                                verify_certs=False,  # 关闭 SSL 证书验证
                                request_timeout=30  # 设置超时时间，防止连接失败
                                )

    def sync_data(self):
        # 获取MongoDB中的数据
        documents = self.collection.find({}, {'_id': 1, 'problem': 1})

        # 同步到ES
        for doc in documents:
            es_doc = TPADocument(
                mongo_id=str(doc['_id']),
                problem=doc['problem']
            )
            es_doc.save()
        return {'status': 'success', 'message': 'Data sync completed'}

    def sync_single_document(self, mongo_id, problem):
        es_doc = TPADocument(
            mongo_id=str(mongo_id),
            problem=problem
        )
        es_doc.save()
        return {'status': 'success', 'message': 'Document synced'}
